# todo 导入数据， 需要从

import datetime
from stock_pool import shang_zheng
print(len(shang_zheng))
stime = '20210120'
etime = '20220314'
stime1 = datetime.datetime.strptime(stime, '%Y%m%d').strftime('%Y-%m-%d')
etime1 = datetime.datetime.strptime(etime, '%Y%m%d').strftime('%Y-%m-%d')
print("len:", len(shang_zheng))


def calc_var(stocks):
    # symbol = df1["symbol"].iloc[0]
    data = get_price(stocks, stime, etime,
                     '1m',
                     #   close
                     ['close'],
                     True,  # 是否跳过停牌
                     "pre",  # 前复权
                     0,  # 天数
                     is_panel=1)

    dfm = data.to_frame().reset_index()
    dfm.rename(columns={'major': 'date', 'minor': 'symbol'}, inplace=True)
    dfm["day"] = dfm["date"].apply(lambda x: x.strftime("%Y-%m-%d"))
    dfm["minute"] = dfm["date"].apply(lambda x: x.strftime("%H:%M"))
    dfm = dfm.sort_values(['symbol', 'date'])
    dfm = dfm.reset_index(drop=True)

    df_var = dfm.groupby(['symbol', "day"])['close'].agg({"close_var": "std", "close_var_f": lambda x: x[:60].std(),
                                                          "close_f": lambda x: x[:60].iloc[-1] if not x.empty else 0,
                                                          "close_d": lambda x: x.iloc[-1] if not x.empty else 0})
    df_var = df_var.reset_index()
    return df_var


cache_df = None
batch = 20
ii = 0
for i in range(len(shang_zheng) // batch + 1):

    print("-------" + str(i))
    if cache_df is None:
        cache_df = calc_var(shang_zheng[i * batch: (i + 1) * batch])
    else:
        cache_df = pd.concat([cache_df, calc_var(shang_zheng[i * batch: (i + 1) * batch])])

cache_df = cache_df.reset_index()
cache_df.to_csv("./shang_zheng_var_v1.csv", index=False)